#!/usr/bin/python
# -*- coding: utf-8 -*-

from collections import Counter
import tensorflow.contrib.keras as kr
import numpy as np
import os
from pandas import read_csv


base_dir = 'data/message'
train_dir = os.path.join(base_dir, 'message.train.txt')     # 训练集合
test_dir = os.path.join(base_dir, 'message.test.txt')       # 测试集合
val_dir = os.path.join(base_dir, 'message.val.txt')         # 验证集合
vocab_dir = os.path.join(base_dir, 'message.vocab.txt')     # 词库

source_dir = 'data/source'
normal_dir = os.path.join(source_dir, 'CleanNormalChatUser.csv')     # 正常消息目录
forbid_dir = os.path.join(source_dir, 'CleanForbidChatUser.csv')     # 正常消息目录


def init_source():
    """按照训练集：验证集：测试集 = 10：1：2比例分割文件"""
    if not os.path.exists(base_dir):
        os.makedirs(base_dir)

    message_normal = read_csv(normal_dir)
    message_forbid = read_csv(forbid_dir)

    data_normal, data_forbid = [], []
    data_train, data_val, data_test = [], [], []

    for message in message_normal['message']:
        data_normal.append('正常消息\t' + message)
    for message in message_forbid['message']:
        data_forbid.append('垃圾消息\t' + message)

    normal_slice1 = int(10 * len(data_normal) / 13)
    normal_slice2 = int(12 * len(data_normal) / 13)
    forbid_slice1 = int(10 * len(data_forbid) / 13)
    forbid_slice2 = int(12 * len(data_forbid) / 13)

    data_train.extend(data_normal[:normal_slice1])
    data_test.extend(data_normal[normal_slice1 : normal_slice2])
    data_val.extend(data_normal[normal_slice2 : ])

    data_train.extend(data_forbid[: forbid_slice1])
    data_test.extend(data_forbid[forbid_slice1 : forbid_slice2])
    data_val.extend(data_forbid[forbid_slice2 : ])

    open_file(train_dir, mode='w').write('\n'.join(data_train) + '\n')
    open_file(test_dir, mode='w').write('\n'.join(data_test) + '\n')
    open_file(val_dir, mode='w').write('\n'.join(data_val) + '\n')


def open_file(filename, mode='r'):
    """
    Commonly used file reader, change this to switch between python2 and python3.
    mode: 'r' or 'w' for read or write
    """
    return open(filename, mode, encoding='utf-8', errors='ignore')

def read_file(filename):
    """读取文件数据"""
    contents, labels = [], []
    with open_file(filename) as f:
        for line in f:
            try:
                label, content = line.strip().split('\t')
                contents.append(list(content))
                labels.append(label)
            except:
                pass
    return contents, labels

def build_vocab(train_dir, vocab_dir, vocab_size=5000):
    """根据训练集构建词汇表，存储"""
    data_train, _ = read_file(train_dir)

    all_data = []
    for content in data_train:
        all_data.extend(content)

    counter = Counter(all_data)
    count_pairs = counter.most_common(vocab_size - 1)
    words, _ = list(zip(*count_pairs))
    # 添加一个 <PAD> 来将所有文本pad为同一长度
    words = ['<PAD>'] + list(words)

    open_file(vocab_dir, mode='w').write('\n'.join(words) + '\n')

def read_vocab(vocab_dir):
    """读取词汇表"""
    words = open_file(vocab_dir).read().strip().split('\n')
    word_to_id = dict(zip(words, range(len(words))))

    return words, word_to_id

def read_category():
    """读取分类目录，固定"""
    categories = ['垃圾消息', '正常消息']
    cat_to_id = dict(zip(categories, range(len(categories))))

    return categories, cat_to_id

def to_words(content, words):
    """将id表示的内容转换为文字"""
    return ''.join(words[x] for x in content)

def process_file(filename, word_to_id, cat_to_id, max_length=600):
    """将文件转换为id表示"""
    contents, labels = read_file(filename)

    data_id, label_id = [], []
    for i in range(len(contents)):
        data_id.append([word_to_id[x] for x in contents[i] if x in word_to_id])
        label_id.append(cat_to_id[labels[i]])

    # 使用keras提供的pad_sequences来将文本pad为固定长度
    x_pad = kr.preprocessing.sequence.pad_sequences(data_id, max_length)
    y_pad = kr.utils.to_categorical(label_id)  # 将标签转换为one-hot表示

    return x_pad, y_pad

def batch_iter(x, y, batch_size=64):
    """生成批次数据"""
    data_len = len(x)
    num_batch = int((data_len - 1) / batch_size) + 1

    indices = np.random.permutation(np.arange(data_len))
    x_shuffle = x[indices]
    y_shuffle = y[indices]

    for i in range(num_batch):
        start_id = i * batch_size
        end_id = min((i + 1) * batch_size, data_len)
        yield x_shuffle[start_id:end_id], y_shuffle[start_id:end_id]
